import json
import os
import shutil

labels = ["pig", "person"]


def labelme_to_yolo_segment(labelme_json_path, name, output_dir):
    with open(labelme_json_path, "r") as f:
        data = json.load(f)

    image_height = data["imageHeight"]
    image_width = data["imageWidth"]
    yolo_label = ""
    for shape in data["shapes"]:
        label = shape["label"]
        if label not in labels:
            continue
        points = shape["points"]

        # 将多边形点转换为 YOLO 格式
        yolo_format_points = []
        for point in points:
            x = point[0] / image_width
            y = point[1] / image_height
            yolo_format_points.append(f"{x:.6f} {y:.6f}")
        cidx = labels.index(label)
        yolo_label += f"{cidx} " + " ".join(yolo_format_points) + "\n"

        # 保存为 YOLO 格式的标签文件
    output_file_path = os.path.join(
        output_dir, name + ".txt"
    )
    with open(output_file_path, "w") as out_f:
        out_f.write(yolo_label + "\n")


def convert_all_json_in_directory(input_dir, output_dir):
    labels_path = os.path.join(output_dir, "labels")
    os.makedirs(labels_path, exist_ok=True)
    images_path = os.path.join(output_dir, "images")
    os.makedirs(images_path, exist_ok=True)

    for root, dirs, files in os.walk(input_dir):
        for filename in files:
            if filename.endswith(".json"):
                p = root.split("\\")
                name = p[-2] + "_" + p[-1]
                labelme_json_path = os.path.join(root, filename)
                labelme_to_yolo_segment(labelme_json_path, name, labels_path)
                src_img = os.path.join(root, filename.replace(".json", ".png"))
                shutil.copy2(src_img, images_path + "\\" + name + ".png")


input_dir = r"/.vscode/pig"
output_dir = r"C:\Users\kang_\Desktop\workspace\pig_pose\yolo\pig_t"
convert_all_json_in_directory(input_dir, output_dir)
